import argparse
from pathlib import Path
from typing import ClassVar
import polars as pl
from pre_base import DATA_DIR, PrepareBase

# URL = "http://dgv.tcag.ca/dgv/docs/GRCh38_hg38_variants_2020-02-25.txt"


class PrepareDGV(PrepareBase):
    name: ClassVar[str] = "DGV"
    header: ClassVar[str] = '##INFO=<ID=DGV,Number=1,Type=String,Description="DGV, accession:chrom:start:end:gain_af:gain_loss">'

    def read_input(self) -> pl.DataFrame:
        dataframe = pl.read_csv(
            self.input_file,
            separator="\t",
            columns=["variantaccession", "chr", "start", "end", "samplesize", "observedgains", "observedlosses"],
            schema_overrides={"chr": str},
        )
        dataframe = (
            dataframe.with_columns(
                pl.format("chr{}", pl.col("chr")).alias("chr"),
                (pl.col("observedgains") / pl.col("samplesize")).alias("gain_af"),
                (pl.col("observedlosses") / pl.col("samplesize")).alias("loss_af"),
            )
            .with_columns(
                pl.format(
                    "{}:{}:{}:{}:{}:{}",
                    pl.col("variantaccession"),
                    pl.col("chr"),
                    pl.col("start"),
                    pl.col("end"),
                    pl.col("gain_af"),
                    pl.col("loss_af"),
                ).alias("INFO")
            )
            .filter(pl.col("INFO").is_not_null())
            .select(["chr", "start", "end", "INFO"])
        ).sort(["chr", "start", "end"])
        return dataframe


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Prepare DGV")
    parser.add_argument("--input_file", "-i", default=DATA_DIR / "raw" / "GRCh38_hg38_variants_2020-02-25.txt", type=Path, help="input file")
    parser.add_argument("--outdir", "-o", default=DATA_DIR / "bed", type=Path, help="output directory")
    parser.set_defaults(func=lambda args: PrepareDGV.model_validate(vars(args)).run())
    args = parser.parse_args()
    args.func(args)
